"""This module hosts major terrain analysis utility functions.

- compute_slope
- compute_aspect
- compute_hillshade
"""
import os
import warnings
from typing import Optional, Union

import numpy as np
import numpy.typing as npt
from osgeo import gdal

from loader.nasa_dem.utils import (
    read_data_from_geoimage,
    get_epsg_from_mgrs_tileID,
)
from loader.nasa_dem.config import NASA_DEM_MGRS_COG_ROOT_PATH
from common.geoimage.raster_dataset import RasterDataset
from common.terrain_analysis.utils import reset_nodata


def compute_slope(
    src_dem: Union[gdal.Dataset, str],
    dst_fpath: Optional[str] = None,
    dst_dtype: npt.DTypeLike = np.float32,
    alg: str = "Horn",
    compute_edges: bool = True,
    slope_format: str = "degree",
) -> RasterDataset:
    """Compute the slope from each pixel of a DEM geoimage.

    Parameters
    ----------
    src_dem: str or gdal.Dataset
        Input DEM geoimage file path or GDAL Dataset object.
    dst_fpath: str, optional
        Output file path to the slope surface result map.
    dst_dtype: numpy.dtype
        Output slope array data type: numpy.float32 or numpy.uint16.
    alg: str
        Algorithm used for computing the slope. "Horn" (default) or "ZevenbergenThorne".
    compute_edges: bool
        Whether to compute values at dem geoimage edges. True (default) or False.
    slope_format: str
        Format of the output slope map. "degree" (default) or "percent".

    Returns
    -------
    (np_arr, meta): Tuple(numpy.ndarray, SceneMeta)
        - np_arr: The slope map of the DEM geoimage.
        - meta: Metadata of the slope map.
    """
    if not isinstance(src_dem, (str, gdal.Dataset)):
        raise TypeError(
            "Invalid type of 'src_dem'. Accepts a 'str' or a 'gdal.Dataset'."
        )
    if isinstance(src_dem, str) and not os.path.exists(src_dem):
        raise FileNotFoundError("Source DEM file does not exists.")
    if dst_dtype not in (np.float32, np.uint16):
        raise ValueError(
            "Invalid value of 'dst_dtype'. Accepts 'np.float32' or 'np.uint16'."
        )
    if alg not in ("Horn", "ZevenbergenThorne"):
        raise ValueError(
            "Invalid value of 'alg'. Accepts 'Horn' or 'ZevenbergenThorne'."
        )
    if slope_format not in ("degree", "percent"):
        raise ValueError(
            "Invalid value of 'slope_format'. Accepts 'degree' or 'percent'."
        )

    if slope_format == "percent" and dst_dtype == np.uint16:
        warnings.warn(
            "The uint16 percent slope may have data overflow for pixels at "
            "steep slope where infinite large percent is possible. "
            "Data type of float32 is recommended for percent slope.",
            category=UserWarning,
        )

    tmp_dst_fpath = "/vsimem/tmp_dem_terrain_analysis.tif"
    if dst_fpath is None:
        dst_fpath = "/vsimem/dst_dem_terrain_analysis.tif"
    try:
        ds = gdal.DEMProcessing(
            tmp_dst_fpath,
            src_dem,
            processing="slope",
            alg=alg,
            computeEdges=compute_edges,
            slopeFormat=slope_format,
        )
    except Exception as e:
        raise Exception(e)
    else:
        if ds is None:
            raise Exception("Returned dataset is none.")
        dst_nodata = np.nan if dst_dtype == np.float32 else 65535
        ds = reset_nodata(ds, dst_nodata=dst_nodata)

        dst_dtype_gdal = gdal.GDT_UInt16 if dst_dtype == np.uint16 else gdal.GDT_Float32
        try:
            ds = gdal.Translate(dst_fpath, ds, outputType=dst_dtype_gdal)
        except Exception as e:
            raise Exception(e)
        else:
            if ds is None:
                raise Exception("Returned dataset is none.")
            np_arr, meta = read_data_from_geoimage(ds)
    finally:
        ds = None
        if gdal.VSIStatL(tmp_dst_fpath):
            gdal.Unlink(tmp_dst_fpath)
        if dst_fpath.startswith("/vsimem/") and gdal.VSIStatL(dst_fpath):
            gdal.Unlink(dst_fpath)

    return RasterDataset.from_ndarray(np_arr, meta)


def compute_aspect(
    src_dem: Union[gdal.Dataset, str],
    dst_fpath: Optional[str] = None,
    dst_dtype: npt.DTypeLike = np.float32,
    alg: str = "Horn",
    compute_edges: bool = True,
    zero_for_flat: bool = False,
) -> RasterDataset:
    """Compute the aspect from each pixel of a DEM geoimage.

    Parameters
    ----------
    src_dem: str or gdal.Dataset
        Input DEM geoimage file path or GDAL Dataset object.
    dst_fpath: str, optional
        Output file path to the aspect surface result map.
    dst_dtype: numpy.dtype
        Output aspect array data type: numpy.float32 or numpy.uint16.
    alg: str
        Algorithm used for computing the slope. 'Horn' (default) or 'ZevenbergenThorne'.
    compute_edges: bool
        Whether to compute values at dem geoimage edges. True (default) or False.
    zero_for_flat: bool
        Whether to return 0 for flat areas with slope=0, instead of -9999. True or False (default).

    Returns
    -------
    (np_arr, meta): Tuple[numpy.ndarray, SceneMeta]
        - np_arr: The aspect map of the DEM geoimage.
        - meta: Metadata of the aspect map.
    """
    if not isinstance(src_dem, (str, gdal.Dataset)):
        raise TypeError(
            "Invalid type of 'src_dem'. Accepts a 'str' or a 'gdal.Dataset'."
        )
    if isinstance(src_dem, str) and not os.path.exists(src_dem):
        raise FileNotFoundError("Source DEM file does not exists.")
    if dst_dtype not in (np.float32, np.uint16):
        raise ValueError(
            "Invalid value of 'dst_dtype'. Accepts 'np.float32' or 'np.uint16'."
        )
    if alg not in ("Horn", "ZevenbergenThorne"):
        raise ValueError(
            "Invalid value of 'alg'. Accepts 'Horn' or 'ZevenbergenThorne'."
        )

    tmp_dst_fpath = "/vsimem/tmp_dem_terrain_analysis.tif"
    if dst_fpath is None:
        dst_fpath = "/vsimem/dst_dem_terrain_analysis.tif"
    try:
        ds = gdal.DEMProcessing(
            tmp_dst_fpath,
            src_dem,
            processing="aspect",
            alg=alg,
            computeEdges=compute_edges,
            zeroForFlat=zero_for_flat,
        )
    except Exception as e:
        raise Exception(e)
    else:
        if ds is None:
            raise Exception("Returned dataset is none.")
        dst_nodata = np.nan if dst_dtype == np.float32 else 65535
        ds = reset_nodata(ds, dst_nodata=dst_nodata)

        dst_dtype_gdal = gdal.GDT_UInt16 if dst_dtype == np.uint16 else gdal.GDT_Float32
        try:
            ds = gdal.Translate(dst_fpath, ds, outputType=dst_dtype_gdal)
        except Exception as e:
            raise Exception(e)
        else:
            if ds is None:
                raise Exception("Returned dataset is none.")
            np_arr, meta = read_data_from_geoimage(ds)
    finally:
        ds = None
        if gdal.VSIStatL(tmp_dst_fpath):
            gdal.Unlink(tmp_dst_fpath)
        if dst_fpath.startswith("/vsimem/") and gdal.VSIStatL(dst_fpath):
            gdal.Unlink(dst_fpath)

    return RasterDataset.from_ndarray(np_arr, meta)


def compute_hillshade(
    src_dem: Union[gdal.Dataset, str],
    sun_azimuth: Union[float, int],
    sun_altitude: Union[float, int],
    dst_fpath: Optional[str] = None,
    alg: str = "Horn",
    compute_edges: bool = True,
) -> RasterDataset:
    """Compute the hillshade from each pixel of a DEM geoimage.

    Parameters
    ----------
    src_dem: str or gdal.Dataset
        Input DEM geoimage file path or GDAL Dataset object.
    sun_azimuth: float
        The azimuth of the sun in [0, 360] degrees.
    sun_altitude: float
        The altitude of the sun in [0, 90] degrees.
    dst_fpath: str, optional
        Output file path to the hillshade surface result map.
    alg: str
        Algorithm used for computing the slope. 'Horn' (default) or 'ZevenbergenThorne'.
    compute_edges: bool
        Whether to compute values at dem geoimage edges. True (default) or False.

    Returns
    -------
    (np_arr, meta): Tuple(numpy.ndarray, SceneMeta)
        - np_arr: The hillshade map of the DEM geoimage.
        - meta: Metadata of the hillshade map.

    Notes
    -----
    To compute hillshade from a DEM raster, the sun azimuth and altitude should be
    obtained first by `utils/get_sun_position.py`
    """
    if not isinstance(src_dem, (str, gdal.Dataset)):
        raise TypeError(
            "Invalid type of 'src_dem'. Accepts a 'str' or a 'gdal.Dataset'."
        )
    if isinstance(src_dem, str) and not os.path.exists(src_dem):
        raise FileNotFoundError("Source DEM file does not exists.")
    if not 0 <= sun_azimuth <= 360:
        raise ValueError(
            "Invalid value of 'sun_azimuth'. Accepts a number in [0, 360] (inclusive)."
        )
    if not 0 <= sun_altitude <= 90:
        raise ValueError(
            "Invalid value of 'sun_altitude'. Accepts a number in [0, 90] (inclusive)."
        )
    if alg not in ("Horn", "ZevenbergenThorne"):
        raise ValueError(
            "Invalid value of 'alg'. Accepts 'Horn' or 'ZevenbergenThorne'."
        )

    if not dst_fpath:
        dst_fpath = "/vsimem/tmp_dem_terrain_analysis.tif"
    try:
        ds = gdal.DEMProcessing(
            dst_fpath,
            src_dem,
            processing="hillshade",
            azimuth=sun_azimuth,
            altitude=sun_altitude,
            alg=alg,
            computeEdges=compute_edges,
        )
    except Exception as e:
        raise Exception(e)
    else:
        if ds is None:
            raise Exception("Returned dataset is none.")
        np_arr, meta = read_data_from_geoimage(ds)
    finally:
        ds = None
        if dst_fpath.startswith("/vsimem/") and gdal.VSIStatL(dst_fpath):
            gdal.Unlink(dst_fpath)

    return RasterDataset.from_ndarray(np_arr, meta)


def quick_test():
    # let us update this to use new catalog later
    from legacy.stats.utils_nasa_dem import retrieve_dem_by_mgrs

    # load dem tile sample for testing
    mgrs_tileID = "15TXJ"
    epsg_code = get_epsg_from_mgrs_tileID(mgrs_tileID)
    geoimage_fname = "NASADEM_MGRS_{}_EPSG_{}.tif".format(mgrs_tileID, epsg_code)
    geoimage_fpath = os.path.join(
        NASA_DEM_MGRS_COG_ROOT_PATH, "T{}".format(mgrs_tileID[:2]), geoimage_fname
    )
    print("geoimage file path: {}".format(geoimage_fpath))

    cur_rst = retrieve_dem_by_mgrs(mgrs_tileID)
    print(cur_rst.meta)

    # compute slope
    # slope_fpath = "/home/pengbo/xbp/terrain_analysis/slope_{}".format(geoimage_fname)
    slope_fpath = None
    cur_rst = compute_slope(
        geoimage_fpath, dst_fpath=slope_fpath, dst_dtype=np.uint16, compute_edges=True
    )
    print(cur_rst.meta)

    # compute aspect
    # aspect_fpath = "/home/pengbo/xbp/terrain_analysis/aspect_{}".format(geoimage_fname)
    aspect_fpath = None
    cur_rst = compute_aspect(
        geoimage_fpath, dst_fpath=aspect_fpath, dst_dtype=np.uint16, compute_edges=True
    )
    print(cur_rst.meta)

    # compute hillshade
    # hillshade_fpath = "/home/pengbo/xbp/terrain_analysis/hillshade_{}".format(
    #     geoimage_fname
    # )
    hillshade_fpath = None
    cur_rst = compute_hillshade(
        geoimage_fpath,
        sun_azimuth=315,
        sun_altitude=45,
        dst_fpath=hillshade_fpath,
        compute_edges=True,
    )
    print(cur_rst.meta)


if __name__ == "__main__":
    quick_test()
